package org.springframework.boot.raft.Election;

import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.*;
import java.util.concurrent.locks.ReentrantLock;

/**
 * Raft算法的实现
 */
public class RaftNode {

    /* ==================== 1. 数据模型 & Key ==================== */

    /**
     * 节点角色
     */
    public enum NodeState {
        FOLLOWER,
        CANDIDATE,
        LEADER
    }

    /**
     * 日志Key：包含 (term, index) 的复合主键
     * 用于在ConcurrentHashMap中保证唯一性。
     */
    public static class LogKey {
        private final long term;
        private final long index;

        public LogKey(long term, long index) {
            this.term = term;
            this.index = index;
        }

        public long getTerm() {
            return term;
        }

        public long getIndex() {
            return index;
        }

        @Override
        public boolean equals(Object o) {
            if (this == o) return true;
            if (o == null || getClass() != o.getClass()) return false;
            LogKey logKey = (LogKey) o;
            return term == logKey.term && index == logKey.index;
        }

        @Override
        public int hashCode() {
            return Objects.hash(term, index);
        }

        @Override
        public String toString() {
            return String.format("(term=%d, index=%d)", term, index);
        }
    }

    /**
     * 日志条目
     */
    public static class LogEntry {
        private final long term;    // 日志所属任期
        private final long index;   // 日志全局顺序索引(从1开始)
        private final String command; // 存储的指令或数据(空字符串表示心跳)

        public LogEntry(long term, long index, String command) {
            this.term = term;
            this.index = index;
            this.command = command;
        }

        public long getTerm() {
            return term;
        }

        public long getIndex() {
            return index;
        }

        public String getCommand() {
            return command;
        }

        @Override
        public String toString() {
            return String.format("LogEntry(term=%d,index=%d,cmd=%s)", term, index, command);
        }
    }

    /* ==================== 2. 节点主要字段 ==================== */

    private final String nodeId;                    // 节点自身ID
    private final AtomicReference<NodeState> state; // 当前节点角色
    private final AtomicLong currentTerm;           // 当前任期
    private final AtomicReference<String> votedFor; // 本任期内投过票给谁

    //使用ConcurrentHashMap来存储日志, LogEntry使用( term,index ) 作为Key来存储日志; value=LogEntry
    private final ConcurrentHashMap<LogKey, LogEntry> logMap;

    // commitIndex: 已经被大多数节点复制的日志索引 (已提交)
    private final AtomicLong commitIndex;
    // lastApplied: 已经应用(写入DB)的最高日志索引
    private final AtomicLong lastApplied;

    // 分配日志时使用的自增索引计数器
    private final AtomicLong lastLogIndexAllocated;

    // ============ 定时线程池 & 选举/心跳定时器 ============
    private final ScheduledExecutorService scheduler;
    private ScheduledFuture<?> electionTimeoutFuture;
    private ScheduledFuture<?> heartbeatFuture; // Leader用于心跳
    private final ReentrantLock lock;

    // 集群中其他节点ID
    private final List<String> peerIds;

    /**
     * 我们新增两个常量, 规定在网络RPC时的最大重试次数，以及在一次选举中尝试投票的最大重选次数
     * (若在一次选举周期内多次发请求Vote仍失败, 就等待下次选举超时重新来过).
     */
    private static final int MAX_RPC_RETRIES = 3;       // 网络调用(投票/AppendEntries)最大重试次数
    private static final int MAX_ELECTION_ROUNDS = 2;   // 在同一个Term中, Candidate多轮发起投票的尝试

    /**
     * 记录当前选举轮次(在Candidate状态下, 同一个Term内若多次投票失败可以自增重试次数)
     */
    private int candidateElectionRetryCount = 0;

    /* ==================== 3. 构造 & 初始化 ==================== */
    public RaftNode(String nodeId, List<String> peerIds) {
        this.nodeId = nodeId;
        this.state = new AtomicReference<>(NodeState.FOLLOWER);
        this.currentTerm = new AtomicLong(0);
        this.votedFor = new AtomicReference<>(null);

        this.logMap = new ConcurrentHashMap<>();
        this.commitIndex = new AtomicLong(0);
        this.lastApplied = new AtomicLong(0);
        this.lastLogIndexAllocated = new AtomicLong(0);

        this.scheduler = Executors.newScheduledThreadPool(2);
        this.lock = new ReentrantLock();
        this.peerIds = (peerIds != null) ? peerIds : List.of();

        // 启动时: FOLLOWER, 设置选举超时
        resetElectionTimeout();
    }

    /* ==================== 4. 随机选举超时 + 开始选举 ==================== */

    /**
     * 重置选举超时：
     * 如果在一定时间内（150~300ms之间随机）没有收到心跳或成为 Leader，则发起选举
     */
    private void resetElectionTimeout() {
        if (electionTimeoutFuture != null) {
            electionTimeoutFuture.cancel(true);
        }
        // 150~300 ms随机
        int timeoutMillis = ThreadLocalRandom.current().nextInt(150, 301);
        //在 timeoutMillis 后，调用 startElection--->即：timeoutMillis时间后，如果没有收到心跳则会进行开始选举
        electionTimeoutFuture = scheduler.schedule(this::startElection, timeoutMillis, TimeUnit.MILLISECONDS);
    }

    /**
     * 开始选举(当Follower超时未收到心跳, 或Candidate超时没成为Leader时调用)
     */
    private void startElection() {
        lock.lock();
        try {
            // 若已是Leader,无需再次发起选举
            if (state.get() == NodeState.LEADER) {
                return;
            }
            // 转成Candidate, 任期+1, 自己投票
            state.set(NodeState.CANDIDATE);
            long newTerm = currentTerm.incrementAndGet();
            votedFor.set(nodeId);

            // 每次新任期开始, candidateElectionRetryCount置0
            candidateElectionRetryCount = 0;

            System.out.println(nodeId + " becomes CANDIDATE in term=" + newTerm);
        } finally {
            lock.unlock();
        }

        requestVotesFromPeers();
    }

    /**
     * 并行发送RequestVote RPC
     * 新增多轮尝试(若一轮投票无果, 再试几轮, 直到超过MAX_ELECTION_ROUNDS就让下次选举超时处理)
     */
    private void requestVotesFromPeers() {
        // 快照当前任期，以防后续更新
        final long termSnapshot = currentTerm.get();

        // 统计投票数，自己先投自己1票
        final AtomicInteger voteCount = new AtomicInteger(1);

        // 倒数锁存器，用于等待所有投票结果
        final CountDownLatch latch = new CountDownLatch(peerIds.size());

        long myLastIndex = getLastLogIndex();
        long myLastTerm = getTermOfLastLog();

        for (String peerId : peerIds) {
            scheduler.execute(() -> {
                //异步执行投票请求,使用多次重试的RPC
                boolean granted = false;
                for (int i = 0; i < MAX_RPC_RETRIES; i++) {
                    // 若成功或失败都跳出
                    if (tryRequestVoteRPC(peerId, termSnapshot, nodeId, myLastIndex, myLastTerm)) {
                        granted = true;
                        break;
                    }
                }
                if (granted) {
                    voteCount.incrementAndGet();
                }
                latch.countDown();
            });
        }

        //阻塞当前线程，最长等待 300ms，直到计数器 latch 变为 0 或超时。目的是等待所有任务完成或超时
        try {
            latch.await(300, TimeUnit.MILLISECONDS);
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
        }

        lock.lock();
        try {
            if (termSnapshot < currentTerm.get()) {
                return; // 本地term被别的节点提高, 这次选举作废
            }
            if (voteCount.get() > peerIds.size() / 2) {
                becomeLeader();
            } else {
                // 未获多数票, 若在本任期内, 允许candidateElectionRetryCount+1再重试一次投票
                candidateElectionRetryCount++;
                if (candidateElectionRetryCount < MAX_ELECTION_ROUNDS) {
                    // 再试一次请求投票(不增加term, 仍是Candidate)
                    System.out.println("Candidate " + nodeId +
                            " no majority, attempt " + candidateElectionRetryCount +
                            "/" + MAX_ELECTION_ROUNDS + " in term " + currentTerm.get());
                    requestVotesFromPeers();
                } else {
                    // 超过多轮尝试, 就等待下一次选举超时
                    System.out.println("Candidate " + nodeId +
                            " gave up this term election. Wait for next timeout...");
                    resetElectionTimeout();
                }
            }
        } finally {
            lock.unlock();
        }
    }

    /**
     * 新增的多次投票RPC封装, 在一个RPC调用里自带失败概率
     */
    private boolean tryRequestVoteRPC(String peerId, long termSnapshot, String candidateId,
                                      long myLastIndex, long myLastTerm) {
        return requestVoteRPC(peerId, termSnapshot, candidateId, myLastIndex, myLastTerm);
    }

    /**
     * 成为Leader后，需要：
     *  1. 取消选举超时
     *  2. 开始周期性发送心跳
     */
    private void becomeLeader() {
        state.set(NodeState.LEADER);
        if (electionTimeoutFuture != null) {
            electionTimeoutFuture.cancel(true);
        }
        System.out.println(nodeId + " becomes LEADER in term=" + currentTerm.get());

        // 启动心跳定时器, 每50ms发送一次空AppendEntries
        startHeartbeat();
    }

    /* ==================== 5. 心跳(空AppendEntries) & 网络重试 ==================== */

    /**
     * Leader 在变更后，启动一个固定速率任务：每隔 50ms 给其他节点发心跳
     */
    private void startHeartbeat() {
        if (heartbeatFuture != null && !heartbeatFuture.isCancelled()) {
            heartbeatFuture.cancel(true);
        }
        heartbeatFuture = scheduler.scheduleAtFixedRate(() -> {
            // 如果我还保持Leader身份，则发送空心跳
            if (state.get() == NodeState.LEADER) {
                sendHeartbeat();
            }
        }, 0, 50, TimeUnit.MILLISECONDS);
    }

    /**
     * 发送空心跳给所有 Follower
     * 这里改进为可多次重试(最多 MAX_RPC_RETRIES 次)
     */
    private void sendHeartbeat() {
        long leaderTerm = currentTerm.get();
        long prevLogIndex = getLastLogIndex();
        long prevLogTerm = getTermOfLastLog();

        for (String peerId : peerIds) {
            // 异步执行
            scheduler.execute(() -> {
                boolean success = false;
                for (int i = 0; i < MAX_RPC_RETRIES; i++) {
                    if (appendEntriesRPC(peerId, leaderTerm, nodeId,
                            prevLogIndex, prevLogTerm,
                            null, null,  // 空日志表示心跳
                            commitIndex.get())) {
                        success = true;
                        break;
                    }
                }
                if (!success) {
                    // 经过多次重试都失败
                    System.out.println("[Heartbeat] Node " + nodeId + " => peer " + peerId
                            + " still fail after " + MAX_RPC_RETRIES + " attempts");
                }
            });
        }
    }

    /* ==================== 6. 投票RPC(模拟调用) & onRequestVote ==================== */

    /**
     * RPC: 请求投票(模拟)
     */
    private boolean requestVoteRPC(String peerId, long term, String candidateId, long lastIndex, long lastTerm) {
        RaftNode peer = SimulationEnv.getNode(peerId);
        if (peer == null) return false;
        return peer.onRequestVote(term, candidateId, lastIndex, lastTerm);
    }

    /**
     * 接收投票请求
     */
    public boolean onRequestVote(long term, String candidateId, long candidateLastIndex, long candidateLastTerm) {
        lock.lock();
        try {
            // 如果对方term更大, 则更新并转Follower
            if (term > currentTerm.get()) {
                currentTerm.set(term);
                state.set(NodeState.FOLLOWER);
                votedFor.set(null);
                if (heartbeatFuture != null) {
                    heartbeatFuture.cancel(true);
                }
                resetElectionTimeout();
            }
            // 如果对方term依旧比我小, 拒绝
            if (term < currentTerm.get()) {
                return false;
            }

            // 若本节点已经投票给别人且不是candidateId,拒绝
            String myVote = votedFor.get();
            if (myVote != null && !myVote.equals(candidateId)) {
                return false;
            }

            // 检查candidate日志是否够新
            long myLastIndex = getLastLogIndex();
            long myLastTerm = getTermOfLastLog();
            boolean upToDate = (candidateLastTerm > myLastTerm)
                    || (candidateLastTerm == myLastTerm && candidateLastIndex >= myLastIndex);
            if (!upToDate) {
                return false;
            }

            // 给对方投票
            votedFor.set(candidateId);
            resetElectionTimeout();
            return true;
        } finally {
            lock.unlock();
        }
    }

    /* ==================== 7. 日志复制 + handleClientRequest + onAppendEntries ==================== */

    /**
     * 客户端请求写日志 -> 仅Leader可执行
     */
    public boolean handleClientRequest(String command) {
        if (state.get() != NodeState.LEADER) {
            System.out.println("Node " + nodeId + " is not LEADER, reject client request:" + command);
            return false;
        }
        lock.lock();
        try {
            long newIndex = lastLogIndexAllocated.incrementAndGet();
            long curTerm = currentTerm.get();
            LogKey key = new LogKey(curTerm, newIndex);
            LogEntry entry = new LogEntry(curTerm, newIndex, command);
            logMap.put(key, entry);

            System.out.println("Leader " + nodeId + " appended log: " + key + " => " + entry);

            // 发送给其他Follower复制
            final AtomicInteger successCount = new AtomicInteger(1);
            final CountDownLatch latch = new CountDownLatch(peerIds.size());
            for (String peerId : peerIds) {
                scheduler.execute(() -> {
                    boolean success = false;
                    for (int i = 0; i < MAX_RPC_RETRIES; i++) {
                        if (appendEntriesRPC(peerId, curTerm, nodeId,
                                getPrevLogIndex(newIndex), getPrevLogTerm(newIndex),
                                key, entry, commitIndex.get())) {
                            success = true;
                            break;
                        }
                    }
                    if (success) {
                        successCount.incrementAndGet();
                    }
                    latch.countDown();
                });
            }

            try {
                latch.await(300, TimeUnit.MILLISECONDS);
            } catch (InterruptedException e) {
                Thread.currentThread().interrupt();
            }

            // 如果过半节点都成功, 提交(更新commitIndex)
            if (successCount.get() > peerIds.size() / 2) {
                long oldCommit = commitIndex.get();
                commitIndex.set(newIndex);
                applyEntries(oldCommit + 1, newIndex);
                return true;
            } else {
                // 没过半成功 -> 回退(模拟只删除这条entry即可)
                logMap.remove(key);
                return false;
            }
        } finally {
            lock.unlock();
        }
    }

    /**
     * Follower接收AppendEntries
     *
     * @param leaderTerm    leader的term
     * @param leaderId      leader节点ID
     * @param prevLogIndex  前一条日志的索引
     * @param prevLogTerm   前一条日志的term
     * @param newLogKey     新日志的key (term+index)，若空表示心跳
     * @param newLogEntry   新日志的entry，若空表示心跳
     * @param leaderCommit  leader的commitIndex
     */
    public boolean onAppendEntries(long leaderTerm, String leaderId,
                                   long prevLogIndex, long prevLogTerm,
                                   LogKey newLogKey, LogEntry newLogEntry,
                                   long leaderCommit) {
        lock.lock();
        try {
            // 1. 检查leaderTerm
            if (leaderTerm > currentTerm.get()) {
                currentTerm.set(leaderTerm);
                state.set(NodeState.FOLLOWER);
                votedFor.set(null);
                // 如果本节点原来在发心跳, 关闭
                if (heartbeatFuture != null) {
                    heartbeatFuture.cancel(true);
                }
                resetElectionTimeout();
            }
            if (leaderTerm < currentTerm.get()) {
                return false;
            }

            // 有效Leader，重置超时
            if (state.get() != NodeState.FOLLOWER) {
                state.set(NodeState.FOLLOWER);
            }
            resetElectionTimeout();

            // 如果是空日志(心跳)，直接返回true即可
            if (newLogKey == null || newLogEntry == null) {
                // Heartbeat AppendEntries, 不需要做日志追加
                // 但是仍要更新commitIndex
                long myCommit = commitIndex.get();
                if (leaderCommit > myCommit) {
                    long newCommit = Math.min(leaderCommit, getLastLogIndex());
                    commitIndex.set(newCommit);
                    applyEntries(myCommit + 1, newCommit);
                }
                return true;
            }

            // 2. 检查prevLogIndex/prevLogTerm是否匹配
            if (prevLogIndex > 0) {
                LogKey prevKey = new LogKey(prevLogTerm, prevLogIndex);
                if (!logMap.containsKey(prevKey)) {
                    // 不匹配 -> 回退
                    removeConflictLogs(prevLogIndex);
                    return false;
                }
            }

            // 3. 追加或覆盖新日志
            LogEntry existing = logMap.get(newLogKey);
            if (existing != null) {
                // 如果已存在相同index但term冲突 -> 删除冲突及之后
                if (existing.getTerm() != newLogEntry.getTerm()) {
                    removeConflictLogs(newLogKey.getIndex());
                    logMap.put(newLogKey, newLogEntry);
                }
            } else {
                // 本地无 -> 直接put
                logMap.put(newLogKey, newLogEntry);
                long allocated = lastLogIndexAllocated.get();
                if (newLogKey.getIndex() > allocated) {
                    // 保证本地lastLogIndexAllocated不小于新来的日志index
                    lastLogIndexAllocated.set(newLogKey.getIndex());
                }
            }

            // 4. 更新commitIndex
            long myCommit = commitIndex.get();
            if (leaderCommit > myCommit) {
                long newCommit = Math.min(leaderCommit, getLastLogIndex());
                commitIndex.set(newCommit);
                applyEntries(myCommit + 1, newCommit);
            }
            return true;
        } finally {
            lock.unlock();
        }
    }

    /**
     * RPC(模拟): Leader发AppendEntries给Follower
     * 带简单网络调用, 保留原10%失败率，但这里会由上层多次重试
     */
    private boolean appendEntriesRPC(String peerId, long leaderTerm, String leaderId,
                                     long prevLogIndex, long prevLogTerm,
                                     LogKey newLogKey, LogEntry newLogEntry,
                                     long leaderCommit) {
        // 模拟环境: 先找到peer node
        RaftNode node = SimulationEnv.getNode(peerId);
        if (node == null) {
            // 网络异常: peer不存在
            return false;
        }

        // 可能有网络故障
        // 简单模拟10%概率调用失败
        if (ThreadLocalRandom.current().nextInt(100) < 10) {
            return false;
        }

        // 若网络正常, 调用对方onAppendEntries
        return node.onAppendEntries(leaderTerm, leaderId, prevLogIndex, prevLogTerm, newLogKey, newLogEntry, leaderCommit);
    }

    /* ==================== 8. 日志冲突回退 & 应用 ==================== */

    /**
     * 用于删除从 conflictIndex 到当前最后的所有日志(可能包含多term)
     */
    private void removeConflictLogs(long conflictIndex) {
        // 遍历logMap, 删除 index >= conflictIndex
        for (LogKey k : logMap.keySet()) {
            if (k.getIndex() >= conflictIndex) {
                logMap.remove(k);
            }
        }
        // 调整 lastLogIndexAllocated
        long maxIndex = getLastLogIndex();
        lastLogIndexAllocated.set(maxIndex);
    }

    /**
     * 当commitIndex前进时, 将[oldCommit+1..newCommit]的日志写入数据库
     */
    private void applyEntries(long startIndex, long endIndex) {
        for (long i = startIndex; i <= endIndex; i++) {
            long index = i;
            Optional<LogKey> maybeKey = logMap.keySet().stream()
                    .filter(k -> k.getIndex() == index)
                    .findFirst();
            if (maybeKey.isPresent()) {
                LogEntry e = logMap.get(maybeKey.get());
                if (e != null) {
                    persistLogEntryToDB(e);
                    lastApplied.incrementAndGet();
                }
            }
        }
    }

    /**
     * 写数据库(示意)
     */
    private void persistLogEntryToDB(LogEntry entry) {
        // 在真实环境中，这里可以执行INSERT/UPDATE语句
        System.out.println("Node " + nodeId + " => Persist to DB: " + entry);
    }

    /* ==================== 9. 工具方法(获取最后日志Index/Term等) ==================== */

    /**
     * 查找本地logMap里 "index" 的最大值
     */
    private long getLastLogIndex() {
        return logMap.keySet().stream().mapToLong(LogKey::getIndex).max().orElse(0);
    }

    /**
     * 获取当前最后日志的term
     */
    private long getTermOfLastLog() {
        long idx = getLastLogIndex();
        if (idx == 0) return 0;
        // 找到 key.index=idx
        Optional<LogKey> maybeKey = logMap.keySet().stream()
                .filter(k -> k.getIndex() == idx)
                .findFirst();
        if (maybeKey.isPresent()) {
            return maybeKey.get().getTerm();
        }
        return 0;
    }

    private long getPrevLogIndex(long newIndex) {
        return (newIndex - 1) == 0 ? 0 : (newIndex - 1);
    }

    private long getPrevLogTerm(long newIndex) {
        long prevIndex = newIndex - 1;
        if (prevIndex <= 0) return 0;
        Optional<LogKey> maybeKey = logMap.keySet().stream()
                .filter(k -> k.getIndex() == prevIndex)
                .findFirst();
        if (maybeKey.isPresent()) {
            return maybeKey.get().getTerm();
        }
        return 0;
    }

    /* ==================== 10. 停止节点 & Getter ==================== */

    public void stopNode() {
        if (electionTimeoutFuture != null) {
            electionTimeoutFuture.cancel(true);
        }
        if (heartbeatFuture != null) {
            heartbeatFuture.cancel(true);
        }
        scheduler.shutdownNow();
        System.out.println("Node " + nodeId + " stopped.");
    }

    public String getNodeId() {
        return nodeId;
    }

    public NodeState getState() {
        return state.get();
    }

    public long getCurrentTerm() {
        return currentTerm.get();
    }

    public long getCommitIndex() {
        return commitIndex.get();
    }

    public long getLastApplied() {
        return lastApplied.get();
    }
}

/* ============== 模拟环境(单机模拟RPC) ============== */
class SimulationEnv {
    private static final ConcurrentHashMap<String, RaftNode> nodeMap = new ConcurrentHashMap<>();

    public static void registerNode(RaftNode node) {
        nodeMap.put(node.getNodeId(), node);
    }

    public static RaftNode getNode(String nodeId) {
        return nodeMap.get(nodeId);
    }

    public static void clear() {
        nodeMap.clear();
    }
}

/* ============== 测试类 ============== */
class RaftTestMain {
    public static void main(String[] args) throws InterruptedException {
        // 1. 创建3个节点
        RaftNode nodeA = new RaftNode("NodeA", List.of("NodeB", "NodeC"));
        RaftNode nodeB = new RaftNode("NodeB", List.of("NodeA", "NodeC"));
        RaftNode nodeC = new RaftNode("NodeC", List.of("NodeA", "NodeB"));

        // 2. 注册
        SimulationEnv.registerNode(nodeA);
        SimulationEnv.registerNode(nodeB);
        SimulationEnv.registerNode(nodeC);

        // 3. 等2秒，看谁成为Leader (注意这里存在10%的模拟网络失败率, 可能多次选举)
        Thread.sleep(2000);

        // 4. 让Leader处理一次客户端请求
        for (RaftNode node : List.of(nodeA, nodeB, nodeC)) {
            if (node.getState() == RaftNode.NodeState.LEADER) {
                boolean result = node.handleClientRequest("CMD-PUT x=y");
                System.out.println(node.getNodeId() + " => handleClientRequest result=" + result);
            }
        }

        // 再等待3秒，观察日志复制情况、心跳
        Thread.sleep(3000);

        // 5. 查看commitIndex / lastApplied
        for (RaftNode node : List.of(nodeA, nodeB, nodeC)) {
            System.out.println(node.getNodeId() + " => state=" + node.getState()
                    + ", term=" + node.getCurrentTerm()
                    + ", commitIndex=" + node.getCommitIndex()
                    + ", lastApplied=" + node.getLastApplied());
        }

        // 6. 停止节点
        nodeA.stopNode();
        nodeB.stopNode();
        nodeC.stopNode();
        SimulationEnv.clear();
    }
}